from typing import Any
import httpx
from a2a.client import A2AClient
from a2a.types import (
    MessageSendParams,
    SendMessageRequest,
    SendStreamingMessageRequest,
)
import asyncio
from uuid import uuid4


send_message_payload: dict[str, Any] = {
    'message': {
        'role': 'user',
        'parts': [{'type': 'text', 'text': 'how much is 10 USD in INR?'}], # 内容对于 Helloworld 智能体来说不重要
        'messageId': uuid4().hex,
    },
}

async def send_message():
    """
    发送非流式消息请求 请求的url是 /message/send
    """
    try:
        async with httpx.AsyncClient() as httpx_client:
            client = await A2AClient.get_client_from_agent_card_url(
                httpx_client=httpx_client, 
                base_url='http://localhost:9999'
            )
            
            request = SendMessageRequest(
                id=uuid4().hex,
                params=MessageSendParams(**send_message_payload)
            )

            response = await client.send_message(request)
            print("非流式响应:")
            print(response.model_dump(mode='json', exclude_none=True))
    except Exception as e:
        print(f"发送非流式消息时出错: {e}")

async def send_streaming_message():
    """
    发送流式消息请求 请求的url是 /message/stream
    """
    try:
        async with httpx.AsyncClient() as httpx_client:
            client = await A2AClient.get_client_from_agent_card_url(
                httpx_client=httpx_client, 
                base_url='http://localhost:9999'
            )
            
            streaming_request = SendStreamingMessageRequest(
                id=uuid4().hex,
                params=MessageSendParams(**send_message_payload)
            )

            # 修正：直接调用，不使用 await
            stream_response = client.send_message_streaming(streaming_request)
            print("流式响应:")
            async for chunk in stream_response:
                print(chunk.model_dump(mode='json', exclude_none=True))
    except Exception as e:
        print(f"发送流式消息时出错: {e}")

if __name__ == '__main__':
    asyncio.run(send_message())
    asyncio.run(send_streaming_message())